import d2lzh_pytorch as d2l
import torch
from torch import nn
import sys
sys.path.append("..")


def corr2d_mutil_in(X, K):
    res = d2l.corr2d(X[0, :, :], K[0, :, :])
    for i in range(1, X.shape[0]):
        res += d2l.corr2d(X[i, :, :], K[i, :, :])
    return res


def corr2d_mutil_out(X, K):
    return torch.stack([corr2d_mutil_in(X, k) for k in K])
